import redis
from redis.commands.search.field import TextField, VectorField
from redis.commands.search.indexDefinition import IndexDefinition
import numpy as np

# 连接到 Redis 实例
redis_client = redis.StrictRedis(
    host='YOUR_REDIS_HOST_NAME',
    port='YOUR_REDIS_PORT',
    password='YOUR_REDIS_ACCESS_KEY'
)

# 创建 Redisearch 索引
index_name = "vector_index"
vector_field_name = "vector_field"
schema = (
    TextField("id"),
    VectorField(vector_field_name, "FLOAT32", 128, "FLAT", {"TYPE": "FLOAT32", "DIM": 128, "DISTANCE_METRIC": "COSINE"})
)

try:
    redis_client.ft(index_name).create_index(schema, definition=IndexDefinition(prefix=["doc:"]))
except Exception as e:
    print(f"Index already exists: {e}")

# 插入向量数据
def insert_vector(redis_client, index_name, doc_id, vector):
    vector_bytes = np.array(vector, dtype=np.float32).tobytes()
    redis_client.hset(f"doc:{doc_id}", mapping={"id": doc_id, vector_field_name: vector_bytes})

# 示例向量数据
doc_id = "1"
vector = np.random.rand(128).tolist()
insert_vector(redis_client, index_name, doc_id, vector)

# 向量搜索
def search_vector(redis_client, index_name, query_vector, k=5):
    query_vector_bytes = np.array(query_vector, dtype=np.float32).tobytes()
    search_client = redis_client.ft(index_name)
    query = f"*=>[KNN {k} @vector_field $vector AS score]"
    params = {"vector": query_vector_bytes}
    results = search_client.search(query, query_params=params).docs
    return results

# 搜索示例
query_vector = np.random.rand(128).tolist()
results = search_vector(redis_client, index_name, query_vector)
for result in results:
    print(f"Found document with ID: {result.id} and score: {result.score}")
